import asyncio
import logging
from typing import Dict, List, Optional, Any, Callable
from enum import Enum, auto
from agents import Agent, RunConfig, Runner, ModelProvider

logger = logging.getLogger(__name__)

class TaskStatus(Enum):
    """任务状态枚举"""
    PENDING = auto()     # 待执行
    RUNNING = auto()     # 执行中
    COMPLETED = auto()   # 已完成
    FAILED = auto()      # 失败
    CANCELLED = auto()   # 已取消

class Task:
    """任务类，用于封装单个任务"""
    def __init__(self,
                 task_id: str,
                 input_data: Any,
                 agent_name: str,
                 priority: int = 0,
                 callback: Optional[Callable] = None):
        self.task_id = task_id
        self.input_data = input_data
        self.agent_name = agent_name
        self.priority = priority
        self.status = TaskStatus.PENDING
        self.result = None
        self.error = None
        self.created_at = asyncio.get_event_loop().time()
        self.started_at = None
        self.completed_at = None
        self.callback = callback

class AgentManager:
    """Agent管理类，用于管理多个Agent实例和任务规划"""
    def __init__(self):
        # 存储已注册的Agent实例
        self.agents: Dict[str, Agent] = {}
        # 存储任务队列
        self.task_queue: List[Task] = []
        # 存储正在运行的任务
        self.running_tasks: Dict[str, asyncio.Task] = {}
        # 存储所有任务
        self.tasks: Dict[str, Task] = {}
        # 模型提供商
        self.model_provider: Optional[ModelProvider] = None
        # 最大并发任务数
        self.max_concurrent_tasks = 5

    def register_agent(self, agent: Agent) -> bool:
        """
        注册一个Agent实例

        Args:
            agent: Agent实例

        Returns:
            bool: 注册是否成功
        """
        try:
            if agent.name in self.agents:
                logger.warning(f"Agent {agent.name} is already registered")
                return False

            self.agents[agent.name] = agent
            logger.info(f"Successfully registered agent: {agent.name}")
            return True
        except Exception as e:
            logger.error(f"Failed to register agent {agent.name}: {str(e)}")
            return False

    def unregister_agent(self, agent_name: str) -> bool:
        """
        注销一个Agent实例

        Args:
            agent_name: Agent名称

        Returns:
            bool: 注销是否成功
        """
        try:
            if agent_name not in self.agents:
                logger.warning(f"Agent {agent_name} is not registered")
                return False

            # 移除Agent
            del self.agents[agent_name]
            logger.info(f"Successfully unregistered agent: {agent_name}")
            return True
        except Exception as e:
            logger.error(f"Failed to unregister agent {agent_name}: {str(e)}")
            return False

    def set_model_provider(self, provider: ModelProvider) -> None:
        """
        设置模型提供商

        Args:
            provider: 模型提供商实例
        """
        self.model_provider = provider

    def create_task(self, task_id: str, input_data: Any, agent_name: str, priority: int = 0, callback: Optional[Callable] = None) -> bool:
        """
        创建一个新任务

        Args:
            task_id: 任务ID
            input_data: 任务输入数据
            agent_name: 执行任务的Agent名称
            priority: 任务优先级(数字越小优先级越高)
            callback: 任务完成后的回调函数

        Returns:
            bool: 创建是否成功
        """
        try:
            # 检查Agent是否已注册
            if agent_name not in self.agents:
                logger.error(f"Agent {agent_name} is not registered")
                return False

            # 创建任务
            task = Task(task_id, input_data, agent_name, priority, callback)
            self.task_queue.append(task)
            # 存储任务到字典
            self.tasks[task_id] = task
            # 按优先级排序任务队列
            self.task_queue.sort(key=lambda x: x.priority)

            logger.info(f"Created task {task_id} for agent {agent_name}")
            return True
        except Exception as e:
            logger.error(f"Failed to create task {task_id}: {str(e)}")
            return False

    async def _execute_task(self, task: Task) -> None:
        """
        执行单个任务

        Args:
            task: 任务实例
        """
        try:
            # 更新任务状态
            task.status = TaskStatus.RUNNING
            task.started_at = asyncio.get_event_loop().time()
            logger.info(f"Starting task {task.task_id} with agent {task.agent_name}")

            # 获取Agent
            agent = self.agents[task.agent_name]

            # 执行任务
            result = await Runner.run(
                starting_agent=agent,
                input=task.input_data,
                run_config=RunConfig(model_provider=self.model_provider)
            )

            # 更新任务结果
            task.status = TaskStatus.COMPLETED
            task.result = result
            task.completed_at = asyncio.get_event_loop().time()
            logger.info(f"Completed task {task.task_id} in {task.completed_at - task.started_at:.2f} seconds")

            # 执行回调
            if task.callback:
                try:
                    task.callback(task)
                except Exception as e:
                    logger.error(f"Callback for task {task.task_id} failed: {str(e)}")

        except Exception as e:
            # 更新任务状态为失败
            task.status = TaskStatus.FAILED
            task.error = str(e)
            task.completed_at = asyncio.get_event_loop().time()
            logger.error(f"Failed to execute task {task.task_id}: {str(e)}")
        finally:
            # 从运行任务列表中移除
            if task.task_id in self.running_tasks:
                del self.running_tasks[task.task_id]

            # 尝试执行下一个任务
            await self._process_queue()

    async def _process_queue(self) -> None:
        """处理任务队列"""
        # 检查是否有空闲槽位
        if len(self.running_tasks) < self.max_concurrent_tasks and self.task_queue:
            # 获取优先级最高的任务
            task = self.task_queue.pop(0)
            # 创建异步任务
            async_task = asyncio.create_task(self._execute_task(task))
            # 存储运行中的任务
            self.running_tasks[task.task_id] = async_task

    async def start(self) -> None:
        """
        启动任务处理
        """
        logger.info("AgentManager started")
        # 初始处理队列
        await self._process_queue()

    async def wait_for_completion(self) -> None:
        """
        等待所有任务完成
        """
        while self.running_tasks or self.task_queue:
            await asyncio.sleep(0.1)

    def get_task_status(self, task_id: str) -> Optional[TaskStatus]:
        """
        获取任务状态

        Args:
            task_id: 任务ID

        Returns:
            TaskStatus: 任务状态，如果任务不存在则返回None
        """
        # 检查运行中的任务
        for task in self.running_tasks.values():
            if hasattr(task, 'task_id') and task.task_id == task_id:
                return task.status

        # 检查任务队列
        for task in self.task_queue:
            if task.task_id == task_id:
                return task.status

        return None

    def get_task_result(self, task_id: str) -> Any:
        """
        获取任务结果

        Args:
            task_id: 任务ID

        Returns:
            Any: 任务结果，如果任务不存在或未完成则返回None
        """
        # 从任务字典中获取任务
        if task_id in self.tasks:
            task = self.tasks[task_id]
            return task.result if task.status == TaskStatus.COMPLETED else None
        return None

# 使用示例
async def main():
    from custom_providers import CustomProvider
    from agents import Agent, ModelSettings

    llm_provider = CustomProvider('deepseek')

    # 创建AgentManager
    manager = AgentManager()
    manager.set_model_provider(llm_provider)

    # 创建并注册Agent
    news_agent = Agent(
        name="NewsCrawlerAgent",
        model_settings=ModelSettings(temperature=0),
        instructions=("你是一个新闻爬虫助手，负责收集和总结新闻信息")
    )
    manager.register_agent(news_agent)

    joke_agent = Agent(
        name="JokeAgent",
        model_settings=ModelSettings(temperature=0.7),
        instructions=("你是一个笑话生成助手，负责创作幽默笑话")
    )
    manager.register_agent(joke_agent)

    # 创建任务
    manager.create_task("task1", "收集今天的科技新闻", "NewsCrawlerAgent", priority=1)
    manager.create_task("task2", "讲个程序员笑话", "JokeAgent", priority=0)
    manager.create_task("task3", "总结本周的财经新闻", "NewsCrawlerAgent", priority=2)

    # 启动任务处理
    await manager.start()
    # 等待所有任务完成
    await manager.wait_for_completion()

    # 获取任务结果
    print(f"Task1 result: {manager.get_task_result('task1')}")
    print(f"Task2 result: {manager.get_task_result('task2')}")
    print(f"Task3 result: {manager.get_task_result('task3')}")

if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(filename)s:%(lineno)d - %(levelname)s - %(message)s')

    asyncio.run(main())